Summary

• Add first-class support for custom data types provided by the Triton repository's triton_kernels module (e.g., Tensor, Storage, StridedLayout).
• Avoid top-level imports for optional triton_kernels.tensor by probing availability and importing lazily at use sites.
• Update CUDA test to emit a single launch event to logs and verify:
  • Exactly one .ndjson log file is produced in the temp logs directory
  • Exactly one launch event exists in that file
  • extracted_args["X"].type == "triton_kernels.tensor.Tensor"

Motivation

• The custom types (Tensor, Storage, StridedLayout) are defined in the Triton repo under triton_kernels and may not be installed by default. Importing them at module import time causes hard failures for users without the extra.
• We still want to reliably validate parsing and logging by asserting that the launch event is recorded and carries the expected structured argument metadata.

Changes

tritonparse/reproducer/utils.py

• Add top-level probe (non-import) for optional dependency:
  • TRITON_KERNELS_CUSTOM_TYPES = importlib.util.find_spec("triton_kernels.tensor") is not None
• Add lazy loader with caching:
  • _get_triton_tensor_types() loads Tensor, Storage, StridedLayout on first use and is decorated with functools.lru_cache(maxsize=1)
• Update _create_arg_from_info branches to:
  • Guard with if not TRITON_KERNELS_CUSTOM_TYPES: raise RuntimeError(...)
  • Import via the cached lazy loader at the point of use
  • Preserve existing tensor dtype handling for random data generation

tests/test_tritonparse.py

• Modify test_triton_kernels_Tensor to:
  • Import from tritonparse.reproducer import utils as reproducer_utils
  • Initialize tritonparse.structured_logging with a temp logs directory
  • Run _topk_forward once to generate a single launch event
  • Read the only .ndjson file under the temp logs directory and assert:
    • Exactly one file
    • Exactly one launch event
    • extracted_args["X"].type == "triton_kernels.tensor.Tensor"
  • Synchronize CUDA and clean up the temp directory (unless TEST_KEEP_OUTPUT=1)

Backward compatibility

• When the optional package is absent, constructing the custom types now raises a clear RuntimeError instructing that the dependency is missing. Callers can catch this or ensure the extra is installed.
• Code paths that do not touch the custom types are unaffected.

Testing

• Unit test updated: TestTritonparseCUDA.test_triton_kernels_Tensor
• How to run (CUDA required):
  • python -m unittest tests.test_tritonparse -v -k test_triton_kernels_Tensor
• Expected: test passes, asserting single launch event and correct X type in extracted_args.